{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e0b884adf3d4f516",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a0472e92df8ee037",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import concurrent.futures\n",
    "import json\n",
    "import re\n",
    "from random import Random\n",
    "\n",
    "from openai import AsyncOpenAI, OpenAI\n",
    "from tqdm import tqdm\n",
    "\n",
    "from virtual_lab.agent import Agent\n",
    "from virtual_lab.constants import CONSISTENT_TEMPERATURE, DEFAULT_FINETUNING_EPOCHS\n",
    "from virtual_lab.run_meeting import run_meeting\n",
    "from virtual_lab.utils import (\n",
    "    async_get_messages,\n",
    "    compute_finetuning_cost,\n",
    "    compute_token_cost,\n",
    "    count_tokens,\n",
    "    get_pubmed_central_article,\n",
    ")\n",
    "\n",
    "from nanobody_constants import (\n",
    "    background_prompt,\n",
    "    nanobody_prompt,\n",
    "    discussions_phase_to_dir,\n",
    "    model as base_model,\n",
    "    model_mini as base_model_mini,\n",
    "    generic_agent,\n",
    "    immunologist,\n",
    "    machine_learning_specialist,\n",
    "    computational_biologist,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df3e8b02a3f2bb9e",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fc28ade3791ee419",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "finetuning_dir = discussions_phase_to_dir[\"finetuning\"]\n",
    "papers_dir = finetuning_dir / \"papers\"\n",
    "summaries_dir = finetuning_dir / \"summaries\"\n",
    "qa_dir = finetuning_dir / \"qa_pairs\"\n",
    "\n",
    "for dir_path in [finetuning_dir, papers_dir, summaries_dir, qa_dir]:\n",
    "    dir_path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "client = OpenAI()\n",
    "async_client = AsyncOpenAI()\n",
    "num_concurrent = 10\n",
    "max_tokens = 250000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "164ca17b6c09fc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Topic to agent mapping\n",
    "topic_to_agent = {\n",
    "    \"nanobodies\": immunologist,\n",
    "    \"SARS-CoV-2 spike protein\": immunologist,\n",
    "    \"SARS-CoV-2 variants KP.3 and JN.1\": immunologist,\n",
    "    \"ESM\": machine_learning_specialist,\n",
    "    \"AlphaFold-Multimer\": computational_biologist,\n",
    "    \"Rosetta\": computational_biologist,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a2085fcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Agent to topics mapping\n",
    "agent_to_topics = {agent: [] for agent in topic_to_agent.values()}\n",
    "\n",
    "for topic, agent in topic_to_agent.items():\n",
    "    agent_to_topics[agent].append(topic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "97d6c25e366d62a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoawait asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "265f6546d2ec031",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def run_query(\n",
    "    semaphore: asyncio.Semaphore, agent: Agent, query: str, model: str = base_model\n",
    ") -> str:\n",
    "    \"\"\"Run a query using the model.\n",
    "\n",
    "    :param semaphore: Semaphore to limit the number of concurrent requests.\n",
    "    :param agent: Agent to use for the query.\n",
    "    :param query: Query to run.\n",
    "    :param model: Model to use for the query.\n",
    "    :return: Response from the model.\n",
    "    \"\"\"\n",
    "    async with semaphore:\n",
    "        assistant = await async_client.beta.assistants.create(\n",
    "            name=agent.title, instructions=agent.prompt, model=model\n",
    "        )\n",
    "        thread = await async_client.beta.threads.create()\n",
    "        await async_client.beta.threads.messages.create(\n",
    "            thread_id=thread.id, role=\"user\", content=query\n",
    "        )\n",
    "        await async_client.beta.threads.runs.create_and_poll(\n",
    "            thread_id=thread.id,\n",
    "            assistant_id=assistant.id,\n",
    "            model=model,\n",
    "            temperature=CONSISTENT_TEMPERATURE,\n",
    "        )\n",
    "        messages = await async_get_messages(client=async_client, thread_id=thread.id)\n",
    "        response = messages[-1][\"content\"][0][\"text\"][\"value\"]\n",
    "\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59f5e2f3a5448780",
   "metadata": {},
   "source": [
    "## PubMed Central search queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb7a07c38d6579a8",
   "metadata": {},
   "source": [
    "Use the agents to generate PubMed Central search queries by topic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ac338d6d173bf0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    concurrent.futures.wait(\n",
    "        [\n",
    "            executor.submit(\n",
    "                run_meeting,\n",
    "                meeting_type=\"individual\",\n",
    "                team_member=agent,\n",
    "                agenda=f'{background_prompt} {nanobody_prompt} You are responsible for understanding the topic \"{topic}\" in the context of designing nanobody binders for SARS-CoV-2. You need to fine-tune yourself on the relevant literature on {topic} to improve your ability to design SARS-CoV-2 nanobody binders. Please write out a series of five distinct search queries that you want to run to find relevant scientific papers on {topic}. Include both queries about {topic} generally as well as queries about how {topic} relates to designing nanobody binders for SARS-CoV-2. Please provide the queries in Python syntax as a list of double-quoted strings.',\n",
    "                agenda_questions=(\n",
    "                    f\"What are the queries that you want to perform to identify the relevant literature on {topic} (as a list of double-quoted strings in Python syntax)?\",\n",
    "                ),\n",
    "                save_dir=finetuning_dir,\n",
    "                save_name=f\"{topic.replace(' ', '_')}_queries\",\n",
    "                temperature=CONSISTENT_TEMPERATURE,\n",
    "            )\n",
    "            for topic, agent in topic_to_agent.items()\n",
    "        ]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4982aa5b11f7381e",
   "metadata": {},
   "source": [
    "Extract the search queries from the agent responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "603675b18435369e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up regex pattern for extracting queries\n",
    "query_list_pattern = re.compile(r'\\[\\s*(\".*?\"\\s*(,\\s*\".*?\"\\s*)*)?,?\\s*\\]')\n",
    "\n",
    "topic_to_queries = {}\n",
    "\n",
    "for topic, agent in topic_to_agent.items():\n",
    "    # Get query path for topic\n",
    "    query_path = finetuning_dir / f\"{topic.replace(' ', '_')}_queries.json\"\n",
    "\n",
    "    # Load query discussion\n",
    "    with open(query_path) as f:\n",
    "        query_discussion = json.load(f)\n",
    "\n",
    "    # Extract queries\n",
    "    query_message = query_discussion[-1][\"message\"]\n",
    "    pattern_result = query_list_pattern.search(query_message)\n",
    "\n",
    "    # Check if pattern is matched\n",
    "    if pattern_result is None:\n",
    "        print(f\"No queries found for {query_path}\")\n",
    "        continue\n",
    "\n",
    "    # Extract queries\n",
    "    queries = json.loads(pattern_result.group())\n",
    "    topic_to_queries[topic] = queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3efba9cdf0a6f0c",
   "metadata": {},
   "source": [
    "## PubMed Central papers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca92296db243ce0c",
   "metadata": {},
   "source": [
    "Have the agents find papers on PubMed Central using the search queries (100 papers per query) and evaluate them based on their abstracts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ba3181c77e2e7b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "for topic, agent in topic_to_agent.items():\n",
    "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "        concurrent.futures.wait(\n",
    "            [\n",
    "                executor.submit(\n",
    "                    run_meeting,\n",
    "                    meeting_type=\"individual\",\n",
    "                    team_member=agent,\n",
    "                    agenda=f'{background_prompt} {nanobody_prompt} You are responsible for understanding the topic \"{topic}\" in the context of designing nanobody binders for SARS-CoV-2. You need to fine-tune yourself on the relevant literature on {topic} to improve your ability to design SARS-CoV-2 nanobody binders. Please use PubMed Central and search for relevant papers on {topic} using the query \"{query}\" and request 100 articles with abstracts only. Read all of the abstracts and based on each abstract individually, decide whether you want to fine-tune yourself on the full text of that paper. Include as many papers as possible, but only include papers that are directly relevant to {topic}. Please provide the PMCIDs and titles of all the papers that you wish to fine-tune yourself on as a Python dictionary mapping PMCID as a double-quoted string to title as a double-quoted string.',\n",
    "                    agenda_questions=(\n",
    "                        \"What are the PMCIDs and titles of the papers you wish to fine-tune yourself on (as a Python dictionary mapping PMCID as a double-quoted string to title as double-quoted string)?\",\n",
    "                    ),\n",
    "                    save_dir=finetuning_dir,\n",
    "                    save_name=f\"{topic.replace(' ', '_')}_papers_{query_num + 1}\",\n",
    "                    temperature=CONSISTENT_TEMPERATURE,\n",
    "                    pubmed_search=True,\n",
    "                )\n",
    "                for query_num, query in enumerate(topic_to_queries[topic])\n",
    "                if not (\n",
    "                    discussions_phase_to_dir[\"finetuning\"]\n",
    "                    / f\"{topic.replace(' ', '_')}_papers_{query_num + 1}.json\"\n",
    "                ).exists()\n",
    "            ]\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47eff2b4e981b180",
   "metadata": {},
   "source": [
    "Extract the selected papers from the agent responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b94895349ed6bb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up regex pattern for extracting queries\n",
    "pmcid_to_title_pattern = re.compile(\n",
    "    r'\\{\\s*(\".*?\"\\s*:\\s*\".*?\"\\s*(,\\s*\".*?\"\\s*:\\s*\".*?\"\\s*)*)?\\}'\n",
    ")\n",
    "\n",
    "for topic, agent in topic_to_agent.items():\n",
    "    # Set up title to PMC ID dictionary\n",
    "    title_to_pmcid = {}\n",
    "    titles_lower, pmcids = set(), set()\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    # Get all paper paths for a topic\n",
    "    paper_paths = sorted(finetuning_dir.glob(f\"{topic_name}_papers_*.json\"))\n",
    "\n",
    "    # Check if all papers results are present\n",
    "    if len(paper_paths) != 5:\n",
    "        print(f\"Missing papers for {topic}\")\n",
    "        continue\n",
    "\n",
    "    # Extract PMC IDs and titles from each papers file\n",
    "    for paper_path in paper_paths:\n",
    "        # Load paper discussion\n",
    "        with open(paper_path) as f:\n",
    "            paper_discussion = json.load(f)\n",
    "\n",
    "        # Extract PMC IDs and titles dictionary\n",
    "        paper_message = paper_discussion[1][\"message\"]\n",
    "        pattern_result = pmcid_to_title_pattern.search(paper_message)\n",
    "\n",
    "        # Check if pattern is matched\n",
    "        if pattern_result is None:\n",
    "            print(f\"No papers found for {paper_path}\")\n",
    "            continue\n",
    "\n",
    "        # Extract PMC IDs and titles dictionary\n",
    "        pmcid_to_title = json.loads(pattern_result.group())\n",
    "\n",
    "        # Add PMC IDs and titles to dictionary, avoiding duplicates\n",
    "        for pmcid, title in pmcid_to_title.items():\n",
    "            # Replace en dash and em dash with a hyphen and convert to lowercase\n",
    "            title = title.replace(\"–\", \"-\").replace(\"—\", \"-\")\n",
    "            title_lower = title.lower()\n",
    "\n",
    "            if title_lower not in titles_lower and pmcid not in pmcids:\n",
    "                title_to_pmcid[title] = pmcid\n",
    "                titles_lower.add(title_lower)\n",
    "                pmcids.add(pmcid)\n",
    "\n",
    "    print(f\"Number of papers found for {topic}: {len(title_to_pmcid):,}\")\n",
    "\n",
    "    # Save title to PMC ID dictionary\n",
    "    with open(finetuning_dir / f\"{topic_name}_title_to_pmcid.json\", \"w\") as f:\n",
    "        json.dump(title_to_pmcid, f, indent=4, sort_keys=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e54a7fd55abb9a8d",
   "metadata": {},
   "source": [
    "Load the PMCIDs and titles from the agent responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "722a42e4e434064d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pmcids = set()\n",
    "\n",
    "for topic in topic_to_agent:\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    with open(finetuning_dir / f\"{topic_name}_title_to_pmcid.json\") as f:\n",
    "        title_to_pmcid = json.load(f)\n",
    "\n",
    "    pmcids.update(title_to_pmcid.values())\n",
    "\n",
    "print(f\"Number of unique PMCIDs: {len(pmcids):,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9749cf90c23f4a08",
   "metadata": {},
   "source": [
    "Download the papers from PubMed Central."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75555d31fe196365",
   "metadata": {},
   "outputs": [],
   "source": [
    "paper_count = 0\n",
    "\n",
    "for pmcid in tqdm(sorted(pmcids)):\n",
    "    title, content = get_pubmed_central_article(pmcid=pmcid)\n",
    "\n",
    "    if title is None:\n",
    "        continue\n",
    "\n",
    "    paper_count += 1\n",
    "\n",
    "    # Save paper\n",
    "    with open(papers_dir / f\"{pmcid}.json\", \"w\") as f:\n",
    "        json.dump({\"title\": title, \"content\": content}, f, indent=4, sort_keys=True)\n",
    "\n",
    "print(f\"Number of papers downloaded: {paper_count:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddc033fce2bae98d",
   "metadata": {},
   "source": [
    "## Summarize papers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40ea1abc921423db",
   "metadata": {},
   "source": [
    "Define a function to using an agent to summarize a paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5b0c48ff62aaf328",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def summarize_paper(\n",
    "    semaphore: asyncio.Semaphore,\n",
    "    agent: Agent,\n",
    "    topic: str,\n",
    "    pmcid: str,\n",
    "    title: str,\n",
    "    content: list[str],\n",
    ") -> tuple[str, str, str]:\n",
    "    \"\"\"Summarize a paper using the model.\n",
    "\n",
    "    :param semaphore: Semaphore to limit the number of concurrent requests.\n",
    "    :param agent: Agent to use for summarization.\n",
    "    :param topic: Topic of interest.\n",
    "    :param pmcid: PMC ID of the paper.\n",
    "    :param title: Title of the paper.\n",
    "    :param content: Content of the paper.\n",
    "    :return: Tuple of PMC ID, title, and summary of the paper.\n",
    "    \"\"\"\n",
    "    # Set up query with paper\n",
    "    query = \"\\n\\n\".join(\n",
    "        [\n",
    "            f'Please summarize in extreme detail the following paper titled \"{title}\". Please focus in particular on summarizing key insights about the topic \"{topic}\" in relation to designing SARS-CoV-2 nanobody binders.'\n",
    "        ]\n",
    "        + content\n",
    "    )\n",
    "\n",
    "    # Run query to get summary\n",
    "    summary = await run_query(\n",
    "        semaphore=semaphore,\n",
    "        agent=agent,\n",
    "        query=query,\n",
    "    )\n",
    "\n",
    "    return pmcid, title, summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cc6e8820de27df6",
   "metadata": {},
   "source": [
    "Use the agents to summarize each paper in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ad321787cb7e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for topic, agent in topic_to_agent.items():\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    # Create save directory\n",
    "    topic_summary_dir = summaries_dir / topic_name\n",
    "    topic_summary_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    # Load title to PMC ID dictionary\n",
    "    with open(finetuning_dir / f\"{topic_name}_title_to_pmcid.json\") as f:\n",
    "        title_to_pmcid: dict[str, str] = json.load(f)\n",
    "\n",
    "    # Get unique PMC IDs\n",
    "    pmcids = sorted(set((title_to_pmcid.values())))\n",
    "\n",
    "    # Load papers\n",
    "    pmcid_to_paper = {}\n",
    "    for pmcid in pmcids:\n",
    "        paper_path = papers_dir / f\"{pmcid}.json\"\n",
    "\n",
    "        if paper_path.exists():\n",
    "            with open(paper_path) as f:\n",
    "                paper: dict[str, str | list[str]] = json.load(f)\n",
    "                pmcid_to_paper[pmcid] = paper\n",
    "\n",
    "    print(f\"Number of papers loaded for {topic}: {len(pmcid_to_paper):,}\")\n",
    "\n",
    "    # Limit papers by length\n",
    "    pmcid_to_paper = {\n",
    "        pmcid: paper\n",
    "        for pmcid, paper in pmcid_to_paper.items()\n",
    "        if sum(count_tokens(paragraph) for paragraph in paper[\"content\"]) <= max_tokens\n",
    "    }\n",
    "\n",
    "    print(\n",
    "        f\"Number of papers after token limit of {max_tokens:,} for {topic}: {len(pmcid_to_paper):,}\"\n",
    "    )\n",
    "\n",
    "    # Compute input token cost\n",
    "    input_token_count = sum(\n",
    "        count_tokens(paragraph)\n",
    "        for paper in pmcid_to_paper.values()\n",
    "        for paragraph in [paper[\"title\"]] + paper[\"content\"]\n",
    "    )\n",
    "\n",
    "    input_token_cost = compute_token_cost(\n",
    "        model=base_model,\n",
    "        input_token_count=input_token_count,\n",
    "        output_token_count=0,\n",
    "    )\n",
    "\n",
    "    print(f\"Approximate input token cost for {topic} = ${input_token_cost:.2f}\")\n",
    "\n",
    "    # Set up semaphore with the number of concurrent requests\n",
    "    semaphore = asyncio.Semaphore(num_concurrent)\n",
    "\n",
    "    # Create tasks for each paper\n",
    "    tasks = [\n",
    "        asyncio.create_task(\n",
    "            summarize_paper(\n",
    "                semaphore=semaphore,\n",
    "                agent=agent,\n",
    "                topic=topic,\n",
    "                pmcid=pmcid,\n",
    "                title=paper[\"title\"],\n",
    "                content=paper[\"content\"],\n",
    "            )\n",
    "        )\n",
    "        for pmcid, paper in pmcid_to_paper.items()\n",
    "    ]\n",
    "\n",
    "    # Run agent summary for each paper\n",
    "    results = [\n",
    "        (await task) for task in tqdm(asyncio.as_completed(tasks), total=len(tasks))\n",
    "    ]\n",
    "\n",
    "    # Save summaries\n",
    "    for pmcid, title, summary in results:\n",
    "        with open(topic_summary_dir / f\"{pmcid}.json\", \"w\") as f:\n",
    "            json.dump(\n",
    "                {\"pmcid\": pmcid, \"title\": title, \"summary\": summary},\n",
    "                f,\n",
    "                indent=4,\n",
    "                sort_keys=True,\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e139b03bfed05c9d",
   "metadata": {},
   "source": [
    "Convert the summaries to the format required for fine-tuning, aggregated by topic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab5352bf603e433",
   "metadata": {},
   "outputs": [],
   "source": [
    "for topic, agent in topic_to_agent.items():\n",
    "    # Convert summaries to training data format\n",
    "    training_data = []\n",
    "\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    # Get summary paths\n",
    "    topic_summary_dir = summaries_dir / topic_name\n",
    "    summary_paths = sorted(topic_summary_dir.glob(\"*.json\"))\n",
    "\n",
    "    for summary_path in summary_paths:\n",
    "        # Load paper summary data\n",
    "        with open(summary_path) as f:\n",
    "            summary_data = json.load(f)\n",
    "\n",
    "        # Extract title and summary\n",
    "        title, summary = summary_data[\"title\"], summary_data[\"summary\"]\n",
    "\n",
    "        # Add example to training data\n",
    "        training_data.append(\n",
    "            {\n",
    "                \"messages\": [\n",
    "                    {\"role\": \"system\", \"content\": agent.prompt},\n",
    "                    {\n",
    "                        \"role\": \"user\",\n",
    "                        \"content\": f'Please tell me about the paper \"{title}\" and its insights into \"{topic}\" in relation to designing SARS-CoV-2 nanobody binders.',\n",
    "                    },\n",
    "                    {\"role\": \"assistant\", \"content\": summary},\n",
    "                ]\n",
    "            }\n",
    "        )\n",
    "\n",
    "    # Shuffle training data\n",
    "    random = Random(0)\n",
    "    random.shuffle(training_data)\n",
    "\n",
    "    # Count tokens\n",
    "    token_count = sum(\n",
    "        count_tokens(message[\"content\"])\n",
    "        for data in training_data\n",
    "        for message in data[\"messages\"]\n",
    "    )\n",
    "\n",
    "    # Determine finetuning cost\n",
    "    finetuning_cost = compute_finetuning_cost(\n",
    "        model=base_model_mini,\n",
    "        token_count=token_count,\n",
    "        num_epochs=DEFAULT_FINETUNING_EPOCHS,\n",
    "    )\n",
    "\n",
    "    # Print stats\n",
    "    print(f\"Number of paper examples for {topic}: {len(training_data):,}\")\n",
    "    print(f\"Token count for {topic}: {token_count:,}\")\n",
    "    print(f\"Finetuning cost for {topic} using {base_model_mini}: ${finetuning_cost:.2f}\")\n",
    "    print()\n",
    "\n",
    "    # Save training data in jsonl format\n",
    "    with open(\n",
    "        finetuning_dir / f\"{topic.replace(' ', '_')}_training_data.jsonl\", \"w\"\n",
    "    ) as f:\n",
    "        f.write(\"\\n\".join(json.dumps(example) for example in training_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a580a12",
   "metadata": {},
   "source": [
    "Convert the summaries to the format required for fine-tuning, aggregated by agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c58fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for agent, topics in agent_to_topics.items():\n",
    "    # Convert summaries to training data format\n",
    "    training_data = []\n",
    "\n",
    "    for topic in topics:\n",
    "        topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "        # Get summary paths\n",
    "        topic_summary_dir = summaries_dir / topic_name\n",
    "        summary_paths = sorted(topic_summary_dir.glob(\"*.json\"))\n",
    "\n",
    "        for summary_path in summary_paths:\n",
    "            # Load paper summary data\n",
    "            with open(summary_path) as f:\n",
    "                summary_data = json.load(f)\n",
    "\n",
    "            # Extract title and summary\n",
    "            title, summary = summary_data[\"title\"], summary_data[\"summary\"]\n",
    "\n",
    "            # Add example to training data\n",
    "            training_data.append(\n",
    "                {\n",
    "                    \"messages\": [\n",
    "                        {\"role\": \"system\", \"content\": agent.prompt},\n",
    "                        {\n",
    "                            \"role\": \"user\",\n",
    "                            \"content\": f'Please tell me about the paper \"{title}\" and its insights into \"{topic}\" in relation to designing SARS-CoV-2 nanobody binders.',\n",
    "                        },\n",
    "                        {\"role\": \"assistant\", \"content\": summary},\n",
    "                    ]\n",
    "                }\n",
    "            )\n",
    "\n",
    "    # Shuffle training data\n",
    "    random = Random(0)\n",
    "    random.shuffle(training_data)\n",
    "\n",
    "    # Count tokens\n",
    "    token_count = sum(\n",
    "        count_tokens(message[\"content\"])\n",
    "        for data in training_data\n",
    "        for message in data[\"messages\"]\n",
    "    )\n",
    "\n",
    "    # Determine finetuning cost\n",
    "    finetuning_cost = compute_finetuning_cost(\n",
    "        model=base_model_mini,\n",
    "        token_count=token_count,\n",
    "        num_epochs=DEFAULT_FINETUNING_EPOCHS,\n",
    "    )\n",
    "\n",
    "    # Print stats\n",
    "    print(f\"Number of paper examples for {agent.title}: {len(training_data):,}\")\n",
    "    print(f\"Token count for {agent.title}: {token_count:,}\")\n",
    "    print(f\"Finetuning cost for {agent.title} using {base_model_mini}: ${finetuning_cost:.2f}\")\n",
    "    print()\n",
    "\n",
    "    # Save training data in jsonl format\n",
    "    with open(\n",
    "        finetuning_dir / f\"{agent.title.replace(' ', '_')}_training_data.jsonl\", \"w\"\n",
    "    ) as f:\n",
    "        f.write(\"\\n\".join(json.dumps(example) for example in training_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29719237",
   "metadata": {},
   "source": [
    "Convert the summaries to the format required for fine-tuning, aggregated across all topics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2306fd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert summaries to training data format\n",
    "training_data = []\n",
    "\n",
    "for topic in topic_to_agent:\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    # Get summary paths\n",
    "    topic_summary_dir = summaries_dir / topic_name\n",
    "    summary_paths = sorted(topic_summary_dir.glob(\"*.json\"))\n",
    "\n",
    "    for summary_path in summary_paths:\n",
    "        # Load paper summary data\n",
    "        with open(summary_path) as f:\n",
    "            summary_data = json.load(f)\n",
    "\n",
    "        # Extract title and summary\n",
    "        title, summary = summary_data[\"title\"], summary_data[\"summary\"]\n",
    "\n",
    "        # Add example to training data\n",
    "        training_data.append(\n",
    "            {\n",
    "                \"messages\": [\n",
    "                    {\"role\": \"system\", \"content\": generic_agent.prompt},\n",
    "                    {\n",
    "                        \"role\": \"user\",\n",
    "                        \"content\": f'Please tell me about the paper \"{title}\" and its insights into \"{topic}\" in relation to designing SARS-CoV-2 nanobody binders.',\n",
    "                    },\n",
    "                    {\"role\": \"assistant\", \"content\": summary},\n",
    "                ]\n",
    "            }\n",
    "        )\n",
    "\n",
    "# Shuffle training data\n",
    "random = Random(0)\n",
    "random.shuffle(training_data)\n",
    "\n",
    "# Count tokens\n",
    "token_count = sum(\n",
    "    count_tokens(message[\"content\"])\n",
    "    for data in training_data\n",
    "    for message in data[\"messages\"]\n",
    ")\n",
    "\n",
    "# Determine finetuning cost\n",
    "finetuning_cost = compute_finetuning_cost(\n",
    "    model=base_model_mini,\n",
    "    token_count=token_count,\n",
    "    num_epochs=DEFAULT_FINETUNING_EPOCHS,\n",
    ")\n",
    "\n",
    "# Print stats\n",
    "print(f\"Number of paper examples: {len(training_data):,}\")\n",
    "print(f\"Token count: {token_count:,}\")\n",
    "print(f\"Finetuning cost using {base_model_mini}: ${finetuning_cost:.2f}\")\n",
    "print()\n",
    "\n",
    "# Save training data in jsonl format\n",
    "with open(\n",
    "    finetuning_dir / \"all_generic_training_data.jsonl\", \"w\"\n",
    ") as f:\n",
    "    f.write(\"\\n\".join(json.dumps(example) for example in training_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b57b1c784cecb01",
   "metadata": {},
   "source": [
    "## Fine-tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45f6e0742afda16f",
   "metadata": {},
   "source": [
    "Upload topic summaries as fine-tuning data and save a mapping from topic to data ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ee5c062a1cd23206",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_to_id = {}\n",
    "\n",
    "for topic in topic_to_agent:\n",
    "    path = finetuning_dir / f\"{topic.replace(' ', '_')}_training_data.jsonl\"\n",
    "\n",
    "    file_object = client.files.create(file=open(path, \"rb\"), purpose=\"fine-tune\")\n",
    "\n",
    "    topic_to_id[topic] = file_object.id\n",
    "\n",
    "with open(finetuning_dir / \"topic_to_id.json\", \"w\") as f:\n",
    "    json.dump(topic_to_id, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91ef1d86",
   "metadata": {},
   "source": [
    "Upload agent summaries as fine-tuning data and save a mapping from agent to data ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba916c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_title_to_id = {}\n",
    "\n",
    "for agent in agent_to_topics:\n",
    "    path = finetuning_dir / f\"{agent.title.replace(' ', '_')}_training_data.jsonl\"\n",
    "\n",
    "    file_object = client.files.create(file=open(path, \"rb\"), purpose=\"fine-tune\")\n",
    "\n",
    "    agent_title_to_id[agent.title] = file_object.id\n",
    "\n",
    "with open(finetuning_dir / \"agent_title_to_id.json\", \"w\") as f:\n",
    "    json.dump(agent_title_to_id, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edfd2dbc",
   "metadata": {},
   "source": [
    "Upload generic summaries as fine-tuning data and save a mapping from generic agent to data ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "68375cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_generic_to_id = {}\n",
    "\n",
    "path = finetuning_dir / \"all_generic_training_data.jsonl\"\n",
    "\n",
    "file_object = client.files.create(file=open(path, \"rb\"), purpose=\"fine-tune\")\n",
    "\n",
    "all_generic_to_id[\"all_generic\"] = file_object.id\n",
    "\n",
    "with open(finetuning_dir / \"all_generic_to_id.json\", \"w\") as f:\n",
    "    json.dump(all_generic_to_id, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b62f6f0d2c26b7f9",
   "metadata": {},
   "source": [
    "Load mapping from topic to data ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bb0e6da655cb759a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(finetuning_dir / \"topic_to_id.json\") as f:\n",
    "    topic_to_id = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88b48e85",
   "metadata": {},
   "source": [
    "Load mapping from agent to data ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8b0d59a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(finetuning_dir / \"agent_title_to_id.json\") as f:\n",
    "    agent_title_to_id = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbbe1d15",
   "metadata": {},
   "source": [
    "Load mapping from generic to data ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1cc2b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(finetuning_dir / \"all_generic_to_id.json\") as f:\n",
    "    all_generic_to_id = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32f052905ad0ce88",
   "metadata": {},
   "source": [
    "Launch fine-tuning jobs by topic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "7da9369b7c17cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for model in [base_model, base_model_mini]:\n",
    "    for topic, file_id in topic_to_id.items():\n",
    "        client.fine_tuning.jobs.create(\n",
    "            training_file=file_id,\n",
    "            model=model,\n",
    "            suffix=topic.replace(\" \", \"_\"),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37fc42ce",
   "metadata": {},
   "source": [
    "Launch fine-tuning jobs by agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87d27ba1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for model in [base_model, base_model_mini]:\n",
    "    for agent_title, file_id in agent_title_to_id.items():\n",
    "        client.fine_tuning.jobs.create(\n",
    "            training_file=file_id,\n",
    "            model=model,\n",
    "            suffix=agent_title.replace(\" \", \"_\"),\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6dd48f8",
   "metadata": {},
   "source": [
    "Launch fine-tuning job for generic agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "c4bfb62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for model in [base_model, base_model_mini]:\n",
    "    client.fine_tuning.jobs.create(\n",
    "        training_file=all_generic_to_id[\"all_generic\"],\n",
    "        model=model,\n",
    "        suffix=\"all_generic\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3704a4f311328397",
   "metadata": {},
   "source": [
    "Check fine-tuning job status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66ae2d0e00cb8e3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(list(client.fine_tuning.jobs.list())[1].fine_tuned_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfd9378d12584190",
   "metadata": {},
   "source": [
    "Set up mapping from agent to fine-tuned model ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "295b1b032e55cc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_to_model = {\n",
    "    \"nanobodies\": \"ft:gpt-4o-2024-08-06:personal:nanobodies:Azf4Ybts\",\n",
    "    \"SARS-CoV-2 spike protein\": \"ft:gpt-4o-2024-08-06:personal:sars-cov-2-spike-protein:Azf4VrrH\",\n",
    "    \"SARS-CoV-2 variants KP.3 and JN.1\": \"ft:gpt-4o-2024-08-06:personal:sars-cov-2-variants-kp-3-and-jn-1:Azes3SvG\",\n",
    "    \"ESM\": \"ft:gpt-4o-2024-08-06:personal:esm:AzfGkSRg\",\n",
    "    \"AlphaFold-Multimer\": \"ft:gpt-4o-2024-08-06:personal:alphafold-multimer:Azv64GGF\",\n",
    "    \"Rosetta\": \"ft:gpt-4o-2024-08-06:personal:rosetta:AzuvllCG\",\n",
    "}\n",
    "\n",
    "topic_to_model_mini = {\n",
    "    \"nanobodies\": \"ft:gpt-4o-mini-2024-07-18:personal:nanobodies:AzeZ4afv\",\n",
    "    \"SARS-CoV-2 spike protein\": \"ft:gpt-4o-mini-2024-07-18:personal:sars-cov-2-spike-protein:AzeYQ1YM\",\n",
    "    \"SARS-CoV-2 variants KP.3 and JN.1\": \"ft:gpt-4o-mini-2024-07-18:personal:sars-cov-2-variants-kp-3-and-jn-1:AzeSvPAe\",\n",
    "    \"ESM\": \"ft:gpt-4o-mini-2024-07-18:personal:esm:AzeQwtOb\",\n",
    "    \"AlphaFold-Multimer\": \"ft:gpt-4o-mini-2024-07-18:personal:alphafold-multimer:AzeewRqW\",\n",
    "    \"Rosetta\": \"ft:gpt-4o-mini-2024-07-18:personal:rosetta:Azeb8IiN\",\n",
    "}\n",
    "\n",
    "agent_to_model = {\n",
    "    immunologist: \"ft:gpt-4o-2024-08-06:personal:immunologist:AzXfaYoL\",\n",
    "    machine_learning_specialist: \"ft:gpt-4o-2024-08-06:personal:machine-learning-specialist:AzXIKoIS\",\n",
    "    computational_biologist: \"ft:gpt-4o-2024-08-06:personal:computational-biologist:AzXNlGca\",\n",
    "}\n",
    "\n",
    "agent_to_model_mini = {\n",
    "    immunologist: \"ft:gpt-4o-mini-2024-07-18:personal:immunologist:AzdyC4jP\",\n",
    "    machine_learning_specialist: \"ft:gpt-4o-mini-2024-07-18:personal:machine-learning-specialist:AzdfrwIj\",\n",
    "    computational_biologist: \"ft:gpt-4o-mini-2024-07-18:personal:computational-biologist:AzdlENYb\",\n",
    "}\n",
    "\n",
    "generic_model = \"ft:gpt-4o-2024-08-06:personal:all-generic:B0DBcePr\"\n",
    "generic_model_mini = \"ft:gpt-4o-mini-2024-07-18:personal:all-generic:B0CaviDa\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a8303b0aafe48c",
   "metadata": {},
   "source": [
    "## Evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b60fd2c0b23933",
   "metadata": {},
   "source": [
    "Create a set of question and answer pairs to evaluate the fine-tuned models based on the summaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6dc6497c4af7946e",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def generate_qa_pairs(\n",
    "    semaphore: asyncio.Semaphore, agent: Agent, pmcid: str, summary: str\n",
    ") -> tuple[str, str]:\n",
    "    \"\"\"Generate question and answer pairs based on the summary using the model.\n",
    "\n",
    "    :param semaphore: Semaphore to limit the number of concurrent requests.\n",
    "    :param agent: Agent to use for summarization.\n",
    "    :param pmcid: PMC ID of the paper.\n",
    "    :param summary: Summary of the paper.\n",
    "    :return: Tuple of PMC ID and response (string of question and answer pairs in JSON dictionary form).\n",
    "    \"\"\"\n",
    "    # Set up query with summary\n",
    "    query = f'Please generate five unique questions based on the following summary of a scientific paper. The questions should be designed so that they are specific to the summary rather than general knowledge, i.e., they can only be answered with information in the summary. However, the questions must include enough context so that they can be answered without knowing which specific paper is being referred to (i.e., do not say things like \"in this study\" or \"in the paper\"). For the answers, please write one correct answer and three incorrect answers for each question. Make the incorrect answers as plausible and as similar to the correct answer as possible while still being incorrect so that selecting the correct answer is challenging without specific knowledge of the paper summary. Your response should *only* be a JSON dictionary mapping from question (string) to answers (list of four strings where the first string is the correct answer and the other three are incorrect). Here is the paper summary:\\n\\n{summary}'\n",
    "\n",
    "    # Run query to get questions and answers\n",
    "    response = await run_query(\n",
    "        semaphore=semaphore,\n",
    "        agent=agent,\n",
    "        query=query,\n",
    "    )\n",
    "\n",
    "    return pmcid, response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9a4d835b5e5b242",
   "metadata": {},
   "outputs": [],
   "source": [
    "for topic, agent in topic_to_agent.items():\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    # Create save directory\n",
    "    topic_qa_dir = qa_dir / topic_name\n",
    "    topic_qa_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    # Get summary paths for the topic\n",
    "    summary_paths = sorted((summaries_dir / topic_name).glob(\"*.json\"))\n",
    "\n",
    "    # Load PMCID to summary mapping\n",
    "    pmcid_to_summary = {}\n",
    "    for summary_path in summary_paths:\n",
    "        # Load paper summary data\n",
    "        with open(summary_path) as f:\n",
    "            summary_data = json.load(f)\n",
    "\n",
    "        # Extract PMCID and summary\n",
    "        pmcid, summary = summary_data[\"pmcid\"], summary_data[\"summary\"]\n",
    "\n",
    "        # Add PMCID to summary mapping\n",
    "        pmcid_to_summary[pmcid] = summary\n",
    "\n",
    "    print(f\"Number of summaries loaded for {topic}: {len(pmcid_to_summary):,}\")\n",
    "\n",
    "    # Compute input token cost\n",
    "    input_token_count = sum(\n",
    "        count_tokens(summary) for summary in pmcid_to_summary.values()\n",
    "    )\n",
    "\n",
    "    input_token_cost = compute_token_cost(\n",
    "        model=base_model,\n",
    "        input_token_count=input_token_count,\n",
    "        output_token_count=0,\n",
    "    )\n",
    "\n",
    "    print(f\"Approximate input token cost for {topic} = ${input_token_cost:.2f}\")\n",
    "\n",
    "    # Set up semaphore with the number of concurrent requests\n",
    "    semaphore = asyncio.Semaphore(num_concurrent)\n",
    "\n",
    "    # Create tasks for each summary\n",
    "    tasks = [\n",
    "        asyncio.create_task(\n",
    "            generate_qa_pairs(\n",
    "                semaphore=semaphore, agent=agent, pmcid=pmcid, summary=summary\n",
    "            )\n",
    "        )\n",
    "        for pmcid, summary in pmcid_to_summary.items()\n",
    "    ]\n",
    "\n",
    "    # Run agent questions for each paper\n",
    "    results = [\n",
    "        (await task) for task in tqdm(asyncio.as_completed(tasks), total=len(tasks))\n",
    "    ]\n",
    "\n",
    "    # Parse JSON Q&A pairs\n",
    "    pmcid_to_qa_pairs = {}\n",
    "    for pmcid, qa_pairs in results:\n",
    "        try:\n",
    "            pmcid_to_qa_pairs[pmcid] = json.loads(\n",
    "                qa_pairs.replace(\"```json\", \"\").replace(\"```\", \"\").strip()\n",
    "            )\n",
    "        except json.JSONDecodeError:\n",
    "            print(f\"Failed to parse JSON for {pmcid}\")\n",
    "            print(qa_pairs)\n",
    "\n",
    "    # Save summaries\n",
    "    for pmcid, qa_pairs in pmcid_to_qa_pairs.items():\n",
    "        with open(topic_qa_dir / f\"{pmcid}.json\", \"w\") as f:\n",
    "            json.dump(\n",
    "                {\"pmcid\": pmcid, \"qa_pairs\": qa_pairs}, f, indent=4, sort_keys=True\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ae60be5d078daa0",
   "metadata": {},
   "source": [
    "Function to use an agent to answer a question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bec74be5de5f2464",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def answer_question(semaphore: asyncio.Semaphore, agent: Agent, question: str, answers: list[str], model: str) -> \\\n",
    "tuple[\n",
    "    str, str]:\n",
    "    \"\"\"Answer a question using the model.\n",
    "\n",
    "    :param semaphore: Semaphore to limit the number of concurrent requests.\n",
    "    :param agent: Agent to use for answering the question.\n",
    "    :param question: Question to ask.\n",
    "    :param answers: List of answers to choose from in random order.\n",
    "    :param model: Model to use for answering the question.\n",
    "    :return: Tuple of question and selected answer.\n",
    "    \"\"\"\n",
    "    # Set up query with question and answer options\n",
    "    query = f\"Please answer the following question based on the following list of answers. Only one answer is correct. Please select the correct answer and provide as your response the exact text of the answer and nothing else. Here is the question:\\n\\n{question}\\n\\nHere are the possible answers (choose the correct one):\\n\\n{'\\n\\n'.join(answers)}\"\n",
    "\n",
    "    # Run query to get answer\n",
    "    answer = await run_query(\n",
    "        semaphore=semaphore,\n",
    "        agent=agent,\n",
    "        query=query,\n",
    "        model=model,\n",
    "    )\n",
    "\n",
    "    return question, answer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20894f925bb4a125",
   "metadata": {},
   "source": [
    "Evaluate the fine-tuned models on question and answer pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8f8f524c27aa960",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_num_qa_per_topic = 100\n",
    "\n",
    "for topic, agent in topic_to_agent.items():\n",
    "    topic_name = topic.replace(\" \", \"_\")\n",
    "\n",
    "    # Get Q&A paths for the topic\n",
    "    qa_paths = sorted((qa_dir / topic_name).glob(\"*.json\"))\n",
    "\n",
    "    # Load Q&A pairs\n",
    "    qa_pairs = {}\n",
    "    for qa_path in qa_paths:\n",
    "        # Load paper summary data\n",
    "        with open(qa_path) as f:\n",
    "            qa_data = json.load(f)\n",
    "\n",
    "        # Add Q&A pairs\n",
    "        qa_pairs |= qa_data[\"qa_pairs\"]\n",
    "\n",
    "    print(f\"Number of Q&A pairs loaded for {topic}: {len(qa_pairs):,}\")\n",
    "\n",
    "    # Set up randomness\n",
    "    random = Random(0)\n",
    "\n",
    "    # Limit number of Q&A pairs with random sampling\n",
    "    if len(qa_pairs) > max_num_qa_per_topic:\n",
    "        qa_pairs_list = list(qa_pairs.items())\n",
    "        random.shuffle(qa_pairs_list)\n",
    "        qa_pairs = dict(qa_pairs_list[:max_num_qa_per_topic])\n",
    "        print(\n",
    "            f\"Number of Q&A pairs after random sampling for {topic}: {len(qa_pairs):,}\"\n",
    "        )\n",
    "\n",
    "    # Map questions to correct answers\n",
    "    question_to_correct_answer = {\n",
    "        question: answers[0] for question, answers in qa_pairs.items()\n",
    "    }\n",
    "\n",
    "    # Randomize order of answers\n",
    "    for answers in qa_pairs.values():\n",
    "        random.shuffle(answers)\n",
    "\n",
    "    # Loop over agents/models\n",
    "    for answer_agent, model in [\n",
    "        (generic_agent, base_model),\n",
    "        (generic_agent, base_model_mini),\n",
    "        (generic_agent, generic_model),\n",
    "        (generic_agent, generic_model_mini),\n",
    "        (agent, base_model),\n",
    "        (agent, base_model_mini),\n",
    "        (agent, topic_to_model[topic]),\n",
    "        (agent, topic_to_model_mini[topic]),\n",
    "        (agent, agent_to_model[agent]),\n",
    "        (agent, agent_to_model_mini[agent]),\n",
    "    ]:\n",
    "        print(f\"Evaluating {answer_agent.title} with {model} for {topic}.\")\n",
    "\n",
    "        # Set up semaphore with the number of concurrent requests\n",
    "        semaphore = asyncio.Semaphore(num_concurrent)\n",
    "\n",
    "        # Create tasks for each question\n",
    "        tasks = [\n",
    "            asyncio.create_task(\n",
    "                answer_question(\n",
    "                    semaphore=semaphore,\n",
    "                    agent=answer_agent,\n",
    "                    question=question,\n",
    "                    answers=answers,\n",
    "                    model=model,\n",
    "                )\n",
    "            )\n",
    "            for question, answers in qa_pairs.items()\n",
    "        ]\n",
    "\n",
    "        # Run agent answers for each paper\n",
    "        results = [\n",
    "            (await task) for task in tqdm(asyncio.as_completed(tasks), total=len(tasks))\n",
    "        ]\n",
    "\n",
    "        # Map questions to selected answers\n",
    "        question_to_selected_answer = {question: answer for question, answer in results}\n",
    "\n",
    "        # Compute accuracy\n",
    "        accuracy = sum(\n",
    "            question_to_correct_answer[question]\n",
    "            == question_to_selected_answer[question]\n",
    "            for question in question_to_correct_answer\n",
    "        ) / len(question_to_correct_answer)\n",
    "\n",
    "        print(f\"Accuracy: {accuracy:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77b7c3f7cde6f733",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 5,
 "nbformat_minor": 9
}
